from langchain.embeddings import HuggingFaceEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma  # 一个轻量级向量数据库
from langchain.document_loaders import TextLoader

# 1. 初始化 HuggingFace 嵌入模型
# model_name = "sentence-transformers/all-MiniLM-L6-v2"
model_name = "sentence-transformers/all-MiniLM-L6-v2"
model_kwargs = {'device': 'cpu'}  # 使用CPU，如果有GPU可改为 'cuda'
encode_kwargs = {'normalize_embeddings': False} # 是否对向量进行归一化

embeddings = HuggingFaceEmbeddings(
    model_name=model_name,
    model_kwargs=model_kwargs,
    encode_kwargs=encode_kwargs
)

# 2. 测试一下模型：将文本转换为向量
text = "LangChain是一个强大的LLM应用开发框架。"
query_result = embeddings.embed_query(text)
print(f"文本的向量维度: {len(query_result)}")
print(f"前10个向量值: {query_result[:10]}")

# 3. 一个更实用的例子：创建知识库并进行语义搜索
# 加载文档（这里用文本文件示例，你也可以用PDF、网页等）
loader = TextLoader("state_of_the_union.txt", encoding="utf-8") # 准备一个txt文件
documents = loader.load()

# 分割文本
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

print(f"文档被分割成了 {len(docs)} 个片段")

# 4. 为所有文本片段生成向量并存入向量数据库（这里用Chroma）
db = Chroma.from_documents(docs, embeddings)

# 5. 进行相似度查询
query = "总统主要讨论了哪些经济问题？"
retrieved_docs = db.similarity_search(query)

print(f"\n对于查询: '{query}'")
print("检索到的最相关文档片段:")
for i, doc in enumerate(retrieved_docs):
    print(f"\n--- 片段 {i+1} ---")
    print(doc.page_content[:300] + "...") # 打印前300个字符